from socket import inet_aton, inet_ntoa
from binascii import b2a_hex

DEBUG = True
MODE = 2        # 0 = Mainline, 1 = libtorrent/utorrent, 2 = flexible

############################
# Help Functions
############################
def toint(s):
    return long(b2a_hex(s), 16)

def tobinary16(i):
    return chr((i >> 8) & 0xFF) + chr(i & 0xFF)

############################
# Peers Encoding / decoding
############################
def encodePeers(peers):
    if MODE == 0:           # Mainline
        return [encodePeersMainline(peers)]
    else:                   # Others
        return encodePeersOthers(peers)

def decodePeers(peers):
    if MODE == 0:           # Mainline
        return decodePeersMainline(peers[0])
    elif MODE == 1:         # Others
        return decodePeersOthers(peers)
    elif len(peers) == 1:   # Flexible = Mainline
        return decodePeersMainline(peers[0])
    else:                   # Flexible = others
        return decodePeersOthers(peers)
        
def encodePeer(peer):
    """
    compact peer ip/port for single peer
    """
    if type(peer) is dict:
        msg = inet_aton(peer['host']) + tobinary16(peer['port'])
    else:
        msg = inet_aton(peer[0]) + tobinary16(peer[1])
    return msg
 
def encodePeersMainline(peers):
    """
    compact peer ip/port for list of peers
    Mainline style
    """
    if not peers:
        return ""
    if len(peers) == 1:
        return encodePeer(peers[0])
    if type(peers[0]) is dict:
        msg = ""
        for peer in peers:
            msg += inet_aton(peer['host']) + tobinary16(peer['port'])
    else:
        msg = ""
        for peer in peers:
            msg += inet_aton(peer[0]) + tobinary16(peer[1])
    return msg

def encodePeersOthers(peers):
    """
    compact peer ip/port for list of peers
    libtorrent/uTorrent style
    """
    return [encodePeer(peer) for peer in peers if peer is not None]

def decodePeer(string):
    """
    Return (ip, port) from compact peer ip/port string for single peer
    """
    if len(string) == 6:
        return ((inet_ntoa(string[:4]), int(toint(string[4:6]))), 0, None)
    if len(string) < 6:
        if DEBUG:
            print "DHT: invalid peers string len: %d" % len(string)
        return None
    if len(string) == 18:
        if DEBUG:
            print "DHT: TODO IPV6"
        return None
        
def decodePeersMainline(string):
    """
    Return (ip, port) from compact peer ip/port string for list of peers
    Mainline style
    """
    if not string:
        return []
    
    if len(string) % 6:
        if DEBUG:
            print "DHT: invalid peers string len: ", len(string)
        return []
    
    peers = []
    x = 0
    while len(string) > x:
        peers.append(((inet_ntoa(string[x:x+4]), int(toint(string[x+4:x+6]))), 0, None))
        x += 6
    return peers

def decodePeersOthers(stringList):
    """
    Return (ip, port) from compact peer ip/port string for list of peers
    libtorrent/uTorrent style
    """
    peers = [decodePeer(string) for string in stringList]
    peers = filter(None, peers)
    return peers

############################
# Nodes Encoding / decoding
############################
def encodeNodes(nodes):
    if not nodes:
        return ""
    
    msg = ""
    for node in nodes:
        msg += node['id'] + inet_aton(node['host']) + tobinary16(node['port'])

    return msg

def decodeNodes(string):
    if len(string) % 26 != 0:
        raise Exception("invalid nodes string len: %d" % len(string))

    nodes = []
    x = 0
    while len(string) > x:
        node = {}
        node['id'] = string[x:x+20]
        node['host'] = inet_ntoa(string[x+20:x+24])
        node['port'] = int(toint(string[x+24:x+26])) 
        nodes.append(node)
        x += 26
    return nodes
